#include "CGSolver.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
CGSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::CGSolver() : IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>()
{
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CGSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Mult_ipl(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const
{
    Standard(b, x);
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void CGSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Standard(const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const
{
    assert(this->oper != nullptr);

    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    double * record = this->part_times;
    memset(record, 0.0, sizeof(double) * NUM_KRYLOV_RECORD);

    // 初始化辅助向量：r残差，p搜索方向，Ap为A乘以搜索方向
    par_Vector<idx_t, ksp_t, dof> r(x), u(x), p(x), s(x);
    par_Vector<idx_t, pc_calc_t, dof> * pc_buf_b = nullptr, * pc_buf_x = nullptr;

    if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
        pc_buf_b = new par_Vector<idx_t, pc_calc_t, dof>(b.comm, b.glb_nrows, b.beg_row, b.end_row);
        pc_buf_x = new par_Vector<idx_t, pc_calc_t, dof>(*pc_buf_b);
        pc_buf_b->set_val(0.0);
        pc_buf_x->set_val(0.0);
    }
    r.set_val(0.0); u.set_val(0.0); p.set_val(0.0); s.set_val(0.0);

    ksp_t gamma, eta, alpha, beta;
    ksp_t tmp_loc[4], tmp_glb[4];// 用作局部点积的存储，发送缓冲区和接受缓冲区

    if (this->zero_guess) {
        vec_copy(b, r);
    } else {
        record[OPER] -= wall_time();
        this->oper->Mult(x, r, false);
        record[OPER] += wall_time();
        vec_add(b, (ksp_t)-1.0, r, r);
    }
    
    if (this->prec) {// 有预条件子，则以预条件后的残差M^{-1}*r作为搜索方向 p = M^{-1}*r
        bool zero_guess = true;
        if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
            VecTransPrec(r, *pc_buf_b);// copy in
            if (zero_guess) pc_buf_x->set_val(0.0);
            else
                VecTransPrec(u, *pc_buf_x);// copy in
            record[PREC] -= wall_time();
            this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
            record[PREC] += wall_time();
            VecTransPrec(*pc_buf_x, u);
        } else {
            if (zero_guess) u.set_val(0.0);
            record[PREC] -= wall_time();
            this->prec->Mult(r, u, zero_guess);
            record[PREC] += wall_time();
        }
    } else {// 没有预条件则直接以残差r作为搜索方向
        vec_copy(r, u);
    }

    vec_copy(u, p);
    this->oper->Mult(p, s, false);

    // double ck = vec_dot<idx_t, ksp_t, double>(p, p);
    // printf("(p, p)=%.10e\n", ck);
    // ck = vec_dot<idx_t, ksp_t, double>(s, s);
    // printf("(s, s)=%.10e\n", ck);

    if constexpr (sizeof(ksp_t) < 8) {
        tmp_loc[0] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, r);// γ的局部点积
        tmp_loc[1] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(s, p);// η的局部点积
        tmp_loc[2] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(b, b);
        tmp_loc[3] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
    } else {
        tmp_loc[0] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(u, r);// γ的局部点积
        tmp_loc[1] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(s, p);// η的局部点积
        tmp_loc[2] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(b, b);
        tmp_loc[3] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(r, r);
    }
#ifdef __x86_64__
    if constexpr (std::is_same<ksp_t, __float128>::value)
        MPI_Allreduce(tmp_loc, tmp_glb, 4, MPI_LONG_DOUBLE, MPIU_SUM_FP128, comm);
    else
#endif
        MPI_Allreduce(tmp_loc, tmp_glb, 4, sizeof(ksp_t) == 16 ? MPI_LONG_DOUBLE : (sizeof(ksp_t) == 8 ? MPI_DOUBLE : MPI_FLOAT), MPI_SUM, comm);
    gamma = tmp_glb[0];
    eta   = tmp_glb[1];
    double norm_b, norm_r;
#ifdef __x86_64__
    if constexpr (std::is_same<ksp_t, __float128>::value) {
        norm_b = sqrtq(tmp_glb[2]);
        norm_r = sqrtq(tmp_glb[3]);
    } else {
        norm_b = sqrt(tmp_glb[2]);
        norm_r = sqrt(tmp_glb[3]);
    }
#else
    norm_b = sqrt(tmp_glb[2]);
    norm_r = sqrt(tmp_glb[3]);
#endif
    alpha = gamma / eta;

    int & iter = this->final_iter = 0;
    if (my_pid == 0) printf("iter %4d   ||b|| = %.16e ||r||/||b|| = %.16e\n", iter, (double)norm_b, (double)(norm_r / norm_b));

    iter++;// 执行一次预条件子就算一次迭代
    for ( ; iter < this->max_iter; iter++) {
        vec_add(x,  alpha, p, x);
        vec_add(r, -alpha, s, r);

        // double ck = vec_dot<idx_t, ksp_t, double>(x, x);
        // printf("(x, x) = %.5e\n", ck);
        // ck = vec_dot<idx_t, ksp_t, double>(r, r);
        // printf("(r, r) = %.5e\n", ck);

        if (this->prec) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                VecTransPrec(r, *pc_buf_b);
                if (zero_guess) pc_buf_x->set_val(0.0);
                else
                    VecTransPrec(u, *pc_buf_x);// copy in
                record[PREC] -= wall_time();
                this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                record[PREC] += wall_time();
                VecTransPrec(*pc_buf_x, u);
            } else {
                if (zero_guess) u.set_val(0.0);// 注意这里需要把halo区也一起设置为0
                record[PREC] -= wall_time();
                this->prec->Mult(r, u, zero_guess);
                record[PREC] += wall_time();
            }
        } else {
            vec_copy(r, u);
        }
        // ck = vec_dot<idx_t, ksp_t, double>(u,u);
        // printf("(u, u)=%.10e\n", ck);

        ksp_t old_gamma = gamma;
        if constexpr (sizeof(ksp_t) < 8) {
            tmp_loc[0] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(u, r);// γ的局部点积
            tmp_loc[1] = vec_dot_intra_proc<idx_t, ksp_t, double, dof>(r, r);
        } else {
            tmp_loc[0] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(u, r);// γ的局部点积
            tmp_loc[1] = vec_dot_intra_proc<idx_t, ksp_t, ksp_t, dof>(r, r);
        }
        
#ifdef __x86_64__
        if constexpr (std::is_same<ksp_t, __float128>::value)
            MPI_Allreduce(tmp_loc, tmp_glb, 2, MPI_LONG_DOUBLE, MPIU_SUM_FP128, comm);
        else
#endif
            MPI_Allreduce(tmp_loc, tmp_glb, 2, sizeof(ksp_t) == 16 ? MPI_LONG_DOUBLE : (sizeof(ksp_t) == 8 ? MPI_DOUBLE : MPI_FLOAT), MPI_SUM, comm);
        gamma = tmp_glb[0];
#ifdef __x86_64__
        if constexpr (std::is_same<ksp_t, __float128>::value)
            norm_r = sqrtq(tmp_glb[1]);
        else
#endif
            norm_r = sqrt(tmp_glb[1]);

        if (my_pid == 0) printf("iter %4d   alpha %.16e   ||r||/||b|| = %.16e\n", iter, (double)alpha, (double)(norm_r / norm_b));
        if (norm_r / norm_b <= this->rel_tol || norm_r <= this->abs_tol) {
            this->converged = 1;
            break;
        }

        beta = gamma / old_gamma;

        vec_add(u, beta, p, p);
        record[OPER] -= wall_time();
        this->oper->Mult(p, s, false);
        record[OPER] += wall_time();

        eta = this->Dot(p, s);
        if (eta <= 0.0) {
            ksp_t dd = this->Dot(p, p);
            if (dd > 0.0) if (my_pid == 0) printf("WARNING: PCG: The operator is not positive definite. (Ad, d) = %.5e\n", (double)eta);
            if (eta == 0.0) break;
        }
        alpha = gamma / eta;
    }

    if (pc_buf_x) delete pc_buf_x;
    if (pc_buf_b) delete pc_buf_b;
}

template class CGSolver<int, double, double, double, 1>;
template class CGSolver<int, double, double, double, 2>;
template class CGSolver<int, double, double, double, 3>;
template class CGSolver<int, double, double, double, 4>;

template class CGSolver<int, float , double, double, 1>;
template class CGSolver<int, float , double, double, 2>;
template class CGSolver<int, float , double, double, 3>;
template class CGSolver<int, float , double, double, 4>;

template class CGSolver<int, float , float , double, 1>;
template class CGSolver<int, float , float , double, 2>;
template class CGSolver<int, float , float , double, 3>;
template class CGSolver<int, float , float , double, 4>;

template class CGSolver<int, float , float , float , 1>;
template class CGSolver<int, float , float , float , 2>;
template class CGSolver<int, float , float , float , 3>;
template class CGSolver<int, float , float , float , 4>;
#ifdef USE_FP16
template class CGSolver<int, __fp16, float , double, 1>;
template class CGSolver<int, __fp16, float , double, 2>;
template class CGSolver<int, __fp16, float , double, 3>;
template class CGSolver<int, __fp16, float , double, 4>;

template class CGSolver<int, __fp16, float , float , 1>;
template class CGSolver<int, __fp16, float , float , 2>;
template class CGSolver<int, __fp16, float , float , 3>;
template class CGSolver<int, __fp16, float , float , 4>;
#endif
